PyTorch 2.x 모델 배포 표준 torch.export
1. Eager Mode의 한계와 AOT(Ahead-of-Time) 컴파일의 필요성
1.1 PyTorch의 핵심 철학: 동적 계산 그래프와 Eager Execution
PyTorch의 성공은 Pythonic한 문법과 즉시 실행(Eager Execution) 모델의 유연성에 깊이 뿌리내리고 있다. 이 패러다임은 연구 및 개발 단계에서 사용자가 계산 그래프를 동적으로 구성하고 즉각적으로 결과를 확인할 수 있게 하여, 타의 추종을 불허하는 디버깅 편의성과 높은 생산성을 제공한다. 그러나 이러한 동적성은 프로덕션 환경으로 전환될 때 본질적인 한계를 드러낸다. 각 연산이 실행될 때마다 Python 인터프리터를 거치면서 발생하는 오버헤드는 고성능 추론 환경에서 병목으로 작용하며, Python 런타임 자체에 대한 의존성은 모델의 이식성을 심각하게 저해한다.
1.2 프로덕션 배포를 위한 정적 그래프의 당위성
이러한 한계를 극복하고 모델을 다양한 프로덕션 환경—Python 인터프리터가 없는 서버(Python-less environments), 모바일, 엣지 디바이스 등—에 배포하기 위해서는, 모델의 계산 로직을 정적이고 이식 가능한 형태로 직렬화(serialize)하는 과정이 필수적이다. 정적 그래프는 모델의 전체 계산 흐름을 사전에 파악할 수 있게 하므로, 연산자 퓨전(operator fusion), 커널 최적화 등과 같은 Ahead-of-Time (AOT) 최적화를 적용하여 추론 성능을 극대화할 수 있는 기반을 마련한다.
1.3 torch.export의 등장 배경: torch.jit을 넘어선 PyTorch 2.x의 새로운 패러다임
PyTorch는 초창기부터 torch.jit 모듈을 통해 정적 그래프 생성을 시도해왔다. torch.jit.trace는 예제 입력을 기반으로 연산 흐름을 기록하는 방식으로, 데이터에 따라 제어 흐름이 변하는 모델에 대해 건전성(soundness)을 보장하지 못하는 치명적인 결함을 안고 있었다. 반면, torch.jit.script는 Python 코드의 부분 집합인 TorchScript 언어를 정적으로 분석하여 컴파일하는 방식으로 건전성은 보장했지만, 지원하는 문법의 제약으로 인해 사용자가 기존 코드를 대폭 수정해야 하는 불편함이 있었다.
PyTorch 2.0의 출현과 함께, 이러한 딜레마를 해결할 새로운 기술적 토대가 마련되었다. 핵심 기술인 TorchDynamo는 CPython의 Frame Evaluation API를 활용하여 Python 바이트코드를 직접 분석함으로써, 기존 방식보다 훨씬 광범위하고 정확하게 계산 그래프를 캡처하는 능력을 제공한다. torch.export는 바로 이 TorchDynamo 기술 위에 구축된 차세대 모델 export 솔루션으로, torch.jit의 한계를 극복하고 건전하며(sound) 이식 가능한 단일 그래프 표현을 생성하는 것을 목표로 한다.
torch.export의 등장은 단순한 API의 개선을 넘어, PyTorch의 모델 표현(representation)에 대한 철학적 전환을 시사한다. 과거 torch.jit가 “Python 코드를 정적 그래프로 변환하려는 시도“에 가까웠다면, torch.export는 “임의의 Python 코드로부터 순수한 텐서 계산 그래프를 안전하게 추출하는 것“으로 목표를 재정의했다. torch.jit.script가 컴파일러의 편의를 위해 사용자에게 TorchScript라는 제한된 언어 사용을 강제했던 것과 달리, torch.export는 TorchDynamo를 통해 일반적인 Python 코드를 그대로 실행하면서 그래프를 캡처한다. 이는 컴파일러가 더 많은 책임을 짊어짐으로써 사용자 경험을 극대화하는 방향으로의 전환이다. 이 근본적인 차이는 PyTorch가 연구에서의 유연성(‘Pythonic’)과 프로덕션에서의 견고성(‘Portable Graph’)이라는 두 가치를 모두 충족시키려는 전략적 의도를 명확히 보여주며, torch.export는 이 두 세계를 잇는 가장 진보된 형태의 교량 역할을 수행한다.
2. torch.export의 핵심 아키텍처와 ExportedProgram
2.1 내부 동작 메커니즘: PT2 컴파일러 스택의 삼위일체
torch.export의 강력한 기능은 PyTorch 2.x 컴파일러 스택을 구성하는 세 가지 핵심 기술의 유기적인 협력을 통해 구현된다.
- TorchDynamo:
torch.export의 그래프 캡처 능력의 심장부로, Python 바이트코드를 Just-in-Time(JIT) 방식으로 분석하여torch.fx.Graph형태로 변환하는 그래프 획득(Graph Acquisition) 엔진이다. - AOT Autograd: TorchDynamo가 캡처한 그래프를 받아, 추론에 불필요한 역전파(backward pass) 로직을 제거하고,
torch.add_와 같은 인플레이스(in-place) 연산을 부수 효과(side effect)가 없는 함수형(functional) 버전으로 변환한다. 최종적으로 그래프의 모든 연산자를 PyTorch의 C++ 백엔드 연산자 라이브러리인 ATen의 연산자 집합으로 낮추는(lower) 역할을 수행한다. - Torch FX: 캡처된 계산 그래프를 표현하는 중간 표현(Intermediate Representation, IR) 형식이다. 순수 Python으로 구현되어 있어 그래프에 대한 분석 및 변환(transformation)을 유연하게 수행할 수 있는 환경을 제공한다.
2.2 ExportedProgram 객체 심층 분석
torch.export.export() 함수는 torch.nn.Module과 예제 입력을 받아 ExportedProgram이라는 표준화된 객체를 반환한다. 이 객체는 모델을 배포하는 데 필요한 모든 정보를 담고 있으며, 다음과 같은 핵심 속성으로 구성된다.
graph_module(torch.fx.GraphModule): 모델의 실제 계산 로직을 담고 있는 핵심 FX 그래프다. 모든 연산은 ATen 수준으로 표현되며, 디버깅 편의성을 위해 원본 코드의 스택 트레이스(stacktrace)와 같은 메타데이터를 포함한다.graph_signature(ExportGraphSignature): 모델의 입출력 명세를 정의한다. 사용자 입력(user input), 학습 가능한 매개변수(parameter), 그리고 상태를 저장하는 버퍼(buffer)를 명확히 구분하여, 그래프의 정확한 호출 규약(calling convention)을 확립한다.state_dict: 모델의 가중치(매개변수와 버퍼)를 포함하는 딕셔너리다.graph_module자체는 가중치가 없는(parameter-free) 순수 함수 형태로 표현되며,state_dict와 결합될 때 비로소 완전한 모델을 구성한다.range_constraints: 동적 형상(dynamic shapes)을 사용할 경우, 각 심볼릭 차원(symbolic dimension)이 가질 수 있는 값의 범위와 제약 조건을 기술한다.
2.3 ExportedProgram IR의 3대 원칙: 배포를 위한 계약(Contract)
torch.export가 생성하는 ExportedProgram은 다음과 같은 세 가지 핵심 원칙을 엄격하게 준수하며, 이는 배포 과정에서의 안정성과 예측 가능성을 보장하는 기반이 된다.
- Soundness (건전성):
ExportedProgram은 원본 Python 프로그램의 동작과 의미론적으로 동일함을 보장한다. 이는torch.jit.trace가 데이터 종속 제어 흐름에서 보였던 근본적인 결함을 해결한 것이다.1
export 과정에서 추적된 모든 연산과 형상(shape)에 대한 가정(assumption)이 그래프에 기록되며, 이 가정이 유효할 때만 프로그램의 올바른 실행이 보장된다.
-
Normalized (정규화): 그래프 내에는 Python 고유의 의미론(semantics)이 존재하지 않는다. 모든 서브모듈은 인라인(inline)되어 단일의 평탄화된(flattened) 그래프를 형성하고,
getattr노드 없이 모든 매개변수는 그래프의 입력으로 리프트(lift)된다. 연산자 집합은 순수한 ATen 연산자로 제한된다.1 -
Functional (함수형): 그래프는 부수 효과(side effect)가 없는 순수 함수형으로 표현된다. 즉, 입력 텐서를 직접 수정하는 변형(mutation) 연산이나 텐서의 메모리를 공유하는 별칭(aliasing) 연산이 모두 제거된다. 이는 그래프 분석 및 변환을 훨씬 간단하고 안전하게 만든다.1
이 세 가지 원칙은 단순한 기술적 특징을 넘어, PyTorch 프레임워크와 ExecuTorch, TensorRT, ONNX Runtime과 같은 하위 배포 백엔드 간의 엄격한 인터페이스 계약(Strict Interface Contract) 역할을 수행한다. 과거에는 각 백엔드가 PyTorch 모델을 변환하기 위해 저마다의 복잡한 컨버터를 구현해야 했다. 이 컨버터들은 PyTorch의 인플레이스 연산, 모듈 계층 구조, 다양한 Python 자료구조 등을 모두 처리해야 했기 때문에 매우 불안정하고 유지보수가 어려웠다. torch.export는 이 모든 복잡성을 AOT Autograd 단계에서 한 번에 처리하여, 모든 백엔드가 소비할 수 있는 표준화된 ‘정규형(canonical form)’ IR을 제공한다. 결과적으로 백엔드 개발자들은 복잡한 Python 의미론을 고려할 필요 없이, 잘 정의된 순수 계산 그래프에만 집중하여 최적화를 수행할 수 있게 된다. 이는 컴파일러 설계에서 프론트엔드와 백엔드를 명확히 분리하는 원칙과 일맥상통하며, 전체 생태계의 개발 속도와 안정성을 비약적으로 향상시킨다.
3. IR Dialects 심층 분석: 점진적 강하(Progressive Lowering)의 미학
torch.export를 통해 생성된 ExportedProgram은 다양한 배포 타겟에 맞게 추가적인 변환 과정을 거친다. 이 과정에서 ’Dialect’라는 개념이 핵심적인 역할을 한다. Dialect는 특정 목적을 위해 연산자 집합이나 메타데이터에 추가적인 제약을 가한 Exported IR의 변종(variant)으로, ExecuTorch 컴파일 파이프라인은 여러 Dialect를 순차적으로 거치며 프로그램을 점진적으로 하드웨어에 가까운 형태로 변환한다.
3.1 ATen Dialect: 충실한 프로그램의 첫 번째 표현
- 역할: Eager mode PyTorch 프로그램을 Exported IR 그래프로 변환한 직후의 첫 번째 단계를 나타낸다. 이는 ExecuTorch 컴파일 파이프라인의 공식적인 진입점(entry point)이다.2
- 속성: ATen Dialect는 사용자 프로그램을 최대한 충실하게(faithfully) 표현하는 것을 목표로 한다. 연산자 집합은 약 2000개에 달하는 방대한 functional ATen 연산자와 사용자가 등록한 커스텀 연산자로 구성된다. 이 단계에서는 동적 데이터 타입(dtype), 암시적 타입 변환, 브로드캐스팅 등 PyTorch의 유연한 기능 대부분이 유지되며, functionalization(mutation 제거)이 완료된다.
3.2 Edge Dialect: 엣지 환경을 향한 첫걸음
- 역할: 범용 서버 환경이 아닌, 엣지 디바이스 배포에 유용한 특수화(specialization)를 도입하는 단계다. 특정 하드웨어에 종속되지 않는 일반적인 엣지 최적화를 수행하는 것을 목표로 한다.2
- 주요 변환:
- DType Specialization: 연산자의 데이터 타입을 명시적으로 고정하여, 불필요한 타입 검사 로직을 제거하고 커널 선택을 용이하게 함으로써 최종 바이너리 크기를 줄인다.
- Scalar to Tensor Conversion:
int,float과 같은 모든 스칼라 입력을 0차원 텐서로 변환한다. 이를 통해 그래프 내의 모든 데이터 흐름을 텐서로 통일하여 후속 처리 과정을 단순화한다. - Operator Namespace 변경: 연산자의 네임스페이스가
torch.ops.aten에서executorch.exir.dialects.edge로 변경되어, 이 그래프가 엣지 특화 단계를 거쳤음을 명시적으로 나타낸다.
3.3 Backend Dialect: 특정 하드웨어를 위한 최종 변환
- 역할: 특정 백엔드(예: DSP, NPU)에 종속적인 연산자나 최적화를 도입하는 최종 단계다. 이 단계에서 그래프는 타겟 하드웨어에 맞는(target-aware) 형태로 변환된다. 예를 들어, 특정 백엔드에서 제공하는 하드웨어 가속 커널이나 라이브러리를 직접 호출하는 노드가 그래프에 생성될 수 있다.
ATen에서 Edge, 그리고 Backend로 이어지는 이 Dialect 구조는 현대 컴파일러, 특히 MLIR(Multi-Level Intermediate Representation)에서 널리 사용되는 “점진적 강하(Progressive Lowering)” 설계 철학을 채택한 것이다. 만약 단일 단계로 PyTorch Eager 코드에서 최종 하드웨어 코드로 변환하려 한다면, 컴파일러는 추상화 수준이 다른 수많은 최적화를 동시에 고려해야 하므로 극도로 복잡해진다. torch.export의 다이얼렉트 시스템은 이 복잡한 변환 과정을 여러 개의 관리 가능한 단계로 분리한다. ATen Dialect 수준에서는 Python 의미론 제거와 같은 고수준 최적화에 집중하고, Edge Dialect에서는 엣지 환경에 보편적인 중간 수준 최적화를, Backend Dialect에서는 특정 하드웨어에 맞춘 저수준 최적화를 수행한다. 이러한 책임의 분리는 전체 컴파일러의 모듈성, 재사용성, 확장성을 극대화한다. 예를 들어, 새로운 하드웨어 백엔드를 추가할 때, 기존의 ATen에서 Edge로의 변환 로직은 그대로 재사용하고 Backend Dialect로의 변환 패스만 새로 작성하면 되므로, 생태계 확장에 매우 유리한 구조를 제공한다.
4. 레거시 시스템과의 비교 분석: torch.export vs. torch.jit
torch.export의 혁신성을 이해하기 위해서는 이전 세대 기술인 torch.jit과의 비교가 필수적이다.
4.1 torch.jit.trace의 근본적 결함: Soundness의 부재
torch.jit.trace는 예제 입력을 모델에 통과시키며 실행되는 텐서 연산의 흐름을 순차적으로 기록하는 방식이다. 이 방식은 구현이 간단하지만, 입력 데이터 값에 따라 실행 경로가 바뀌는 제어 흐름(control flow)을 정확히 캡처하지 못하는 근본적인 결함을 가지고 있다. 예를 들어, if x.sum() > 0:과 같은 조건문이 있을 때, trace는 예제 입력 x가 통과한 단 하나의 경로만을 기록한다. 만약 다른 입력이 들어와 조건문의 결과가 바뀌면, trace된 모델은 원본 Python 모델과 다르게 동작하게 된다. 이는 “Unsound” 변환의 전형적인 예시로, 배포된 모델의 신뢰성을 심각하게 훼손한다.
4.2 torch.jit.script의 실용적 한계: 표현력의 제약
torch.jit.script는 이러한 문제를 해결하기 위해 Python 코드를 정적으로 분석하여 TorchScript라는 중간 표현(IR)으로 컴파일하는 방식을 채택했다. 이 방식은 제어 흐름을 포함한 더 넓은 범위의 모델 구조를 캡처할 수 있어 trace보다 건전하다. 하지만 이는 사용자가 TorchScript라는, Python 언어의 제한된 부분 집합(subset) 문법에 맞춰 코드를 작성하거나 수정해야 함을 의미했다. 많은 동적 Python 기능이 지원되지 않았기 때문에, 실용적인 관점에서 기존 PyTorch 코드를 script와 호환되도록 만드는 것은 종종 상당한 노력을 요구했다.
4.3 torch.export의 혁신: TorchDynamo를 통한 건전하고 유연한 캡처
torch.export는 TorchDynamo를 기반으로 Python 바이트코드를 직접 분석함으로써 이 두 가지 레거시 방식의 장점만을 취한다. torch.jit.script처럼 정적 분석을 통해 건전성을 보장하면서도, 실제 Python 코드를 실행하며 그래프를 캡처하기 때문에 훨씬 넓은 범위의 Python 기능을 자연스럽게 지원한다. 그 결과물인 ExportedProgram은 순수 ATen 연산자로 구성된 단일 정규형 그래프이므로, torch.jit의 복잡한 IR보다 후속 백엔드에서의 분석 및 최적화가 훨씬 용이하다.
4.4 PyTorch 모델 직렬화/컴파일 기술 비교
다음 표는 각 기술의 핵심 특징, 장단점, 주요 사용 사례를 요약하여 기술 선택에 대한 명확한 가이드를 제공한다.
| 특징 (Feature) | torch.export | torch.compile | torch.jit.trace | torch.jit.script |
|---|---|---|---|---|
| 컴파일 시점 | AOT (Ahead-of-Time) | JIT (Just-in-Time) | AOT (Tracing-Time) | AOT (Scripting-Time) |
| 그래프 캡처 | 전체 그래프 (단일) | 부분 그래프 (Graph Breaks 허용) | 부분 그래프 (단일 경로) | 전체 그래프 (조건부) |
| Soundness | 보장 (Guaranteed) | N/A (JIT 컨텍스트) | 보장 안 됨 (Unsound) | 보장 (Guaranteed) |
| Python 지원 | 높음 (Dynamo 기반) | 매우 높음 (Fallback 허용) | 제한적 (Tensor 연산만) | 제한적 (TorchScript 언어) |
| 주요 출력물 | ExportedProgram (ATen Dialect) | 최적화된 Python Callable | ScriptModule | ScriptModule |
| 이식성 | 매우 높음 (Python-less) | 낮음 (Python 런타임 의존) | 중간 | 높음 |
| 주요 사용 사례 | 모델 배포 (엣지, 서버) | 학습/추론 가속 | 간단한 모델 배포 (레거시) | 제어 흐름 있는 모델 배포 (레거시) |
| 핵심 기술 | TorchDynamo, AOT Autograd | TorchDynamo, Inductor | 연산 기록 | TorchScript 컴파일러 |
5. 고급 기능 활용: 동적 형상(Dynamic Shapes)과 제어 흐름(Control Flow)
정적 그래프의 이점을 취하면서도 실제 배포 시나리오에서 요구되는 유연성을 확보하기 위해, torch.export는 동적 형상과 데이터 종속 제어 흐름을 처리하는 정교한 메커니즘을 제공한다.
5.1 동적 형상(Dynamic Shapes) 처리 메커니즘
실제 추론 환경에서는 배치 크기나 입력 시퀀스의 길이가 매번 달라질 수 있다. torch.export는 이러한 동적성을 정적 그래프 내에서 표현하기 위해 다음과 같은 체계를 사용한다.
torch.export.Dim객체: 동적일 수 있는 텐서의 차원을 명시적으로 선언하는 데 사용된다.Dim("batch", min=2)와 같이 심볼릭 이름과 함께 해당 차원이 가질 수 있는 값의 최소/최대 범위를 지정하여 컴파일러에게 추가적인 정보를 제공할 수 있다.- 심볼릭 Shape과
range_constraints:Dim으로 선언된 차원은$s_0$,$s_1$과 같은 구체적인 숫자 값이 아닌 심볼(symbol)로 그래프에 표현된다.ExportedProgram의range_constraints속성은 이 심볼들이 가질 수 있는 값의 범위와 다른 심볼과의 관계(예:$s_1 = s_0 * 2$)를 저장한다. - 가드(Guards)의 역할:
torch.export는 추적 과정에서 발견된 형상에 대한 가정(예: 브로드캐스팅을 위한 차원 일치 조건$s_1 = 5$)을 “가드“로 생성한다. 추론 시 입력 텐서가 이 가드 조건을 만족하지 않으면 오류가 발생하여, 컴파일된 그래프의 의미론적 정확성을 보장한다. 이는torch.compile의 가드 모델과 유사한 원리로 동작한다.
아래는 dynamic_shapes 인자를 사용하여 동적 형상을 처리하는 예제 코드다.
import torch
from torch.export import export, Dim
class DynamicModel(torch.nn.Module):
def forward(self, x, y):
return x + y
# 배치 차원을 동적으로 지정
batch = Dim("batch", min=2)
dynamic_shapes = {
"x": (batch, 128),
"y": (batch, 128),
}
model = DynamicModel()
example_args = (torch.randn(4, 128), torch.randn(4, 128))
# 동적 형상 정보와 함께 export
ep = export(model, example_args, dynamic_shapes=dynamic_shapes)
# ep.graph_module을 출력하면 x와 y의 첫 번째 차원이
# 구체적인 숫자 '4'가 아닌 심볼 's0'으로 표현된 것을 확인할 수 있다.
# ep.range_constraints를 출력하면 's0'의 범위가.
torch.cond: 데이터 종속적인 조건 분기(if/else)를 표현한다. 불리언 값을 반환하는predicate함수와, 조건이 참일 때와 거짓일 때 각각 실행될true_fn,false_fn을 인자로 받는다.cond는 몇 가지 제약 조건(예: 분기 함수의 입출력은 텐서여야 함)을 가지며, 이를 통해 제어 흐름을 정적으로 분석 가능한 형태로 변환한다.torch.map: 텐서의 특정 차원(주로 첫 번째 차원)을 따라 함수를 반복적으로 적용하는 연산을 표현한다. 이는for루프와 유사한 기능을 수행하며, 루프를 병렬화 가능한 형태로 추상화한다.
6. 한계점 인식과 해결 방안
torch.export는 강력한 도구이지만, 모든 PyTorch 및 Python 코드를 완벽하게 처리할 수는 없다. 그 한계를 명확히 인식하고 해결 방안을 아는 것이 중요하다.
6.1 그래프 분기(Graph Breaks): torch.compile과의 결정적 차이
torch.export의 가장 중요하고 근본적인 제약은 그래프 분기(Graph Breaks)를 허용하지 않는다는 점이다. torch.compile은 추적 불가능한 Python 코드를 만나면 해당 부분의 실행을 Python 인터프리터에게 위임하고(폴백, fallback) 나머지 부분에 대해서만 컴파일을 계속한다. 반면, torch.export는 이러한 상황에서 에러를 발생시킨다. 이는 torch.export의 핵심 목표가 Python 런타임에 대한 의존성이 전혀 없는, 완전히 독립적이고 이식 가능한 단일 그래프를 생성하는 것이기 때문이다. 그래프 분기는 이 목표에 정면으로 위배되는 행위이므로, 의도적으로 금지된 설계 선택이다.
6.2 ExportDB: 호환성 문제의 표준 레퍼런스
torch.export가 실패했을 때, 사용자는 ExportDB를 통해 문제의 원인을 파악하고 해결의 실마리를 찾을 수 있다. ExportDB는 torch.export에서 지원되거나 지원되지 않는 다양한 Python 및 PyTorch 기능에 대한 코드 예제를 모아놓은 공식 데이터베이스다.3 사용자는 모델을 export하기 전에 ExportDB를 참조하여 잠재적인 호환성 문제를 미리 확인하거나, export 실패 시 유사한 실패 사례와 권장되는 코드 수정 방안을 찾아볼 수 있다.
6.3 우회 전략 및 고급 기법
추적 불가능한 코드를 만났을 때 다음과 같은 전략을 고려할 수 있다.
- Non-Strict Export (
strict=False):export함수에strict=False옵션을 전달하면, 지원되지 않는 Python 코드를 만났을 때 에러를 발생시키는 대신 해당 연산을 그래프에 기록하지 않고 조용히 건너뛴다. 예를 들어, 텐서 연산이 아닌id(x)와 같은 함수 호출은 무시된다. 이는 일단 export를 성공시키는 데 도움이 될 수 있지만, 결과 그래프가 원본 프로그램과 다르게 동작할 수 있는 위험이 있으므로 매우 신중하게 사용해야 한다. - 사용자 정의 연산자(Custom Operators): PyTorch가 기본적으로 지원하지 않는 연산을
torch.libraryAPI를 통해 명시적으로 등록하여torch.export가 인식하고 그래프에 포함시킬 수 있다. 이는 C++나 CUDA로 구현된 고성용 커널을torch.export워크플로우에 완벽하게 통합하는 가장 올바르고 강력한 방법이다.4
“No Graph Breaks” 원칙과 ExportDB의 존재는 torch.export가 사용자에게 **“명시성(Explicitness)”**을 요구하고 있음을 보여준다. torch.compile이 암시적인 폴백을 통해 “어떻게든 작동하게” 만드는 것을 목표로 한다면, torch.export는 모델의 모든 계산 흐름이 명시적으로 그래프에 표현될 수 있도록 사용자의 적극적인 개입과 코드 수정을 유도한다. 배포된 모델은 어떤 입력에 대해서도 항상 예측 가능하고 동일하게 동작해야 한다. 만약 모델의 일부가 Python 인터프리터에 의해 실행된다면, 이는 숨겨진 런타임 의존성을 만들고 예측 가능성을 해친다. torch.export는 이 문제를 원천 차단하기 위해 그래프 분기를 금지하고, 대신 ExportDB, torch.cond, torch.library와 같은 도구를 제공하여 사용자가 자신의 의도를 컴파일러가 이해할 수 있는 명시적인 형태로 변환하도록 돕는다. 따라서 torch.export 사용 시 발생하는 에러는 단순한 버그가 아니라, “코드에 컴파일러가 이해할 수 없는 모호한 부분이 있으니 명확하게 수정해달라“는 컴파일러의 요구사항으로 해석해야 한다. 이는 장기적으로 훨씬 더 견고하고 신뢰할 수 있는 모델 배포로 이어진다.
7. 배포 생태계 통합 워크플로우
torch.export를 통해 생성된 ExportedProgram은 다양한 배포 백엔드로 향하는 표준화된 시작점 역할을 한다.
7.1 Executorch: 엣지 AI를 위한 표준 경로
torch.export는 엣지 디바이스용 추론 프레임워크인 ExecuTorch 파이프라인의 공식적인 시작점이다. 전체 워크플로우는 다음과 같다.
torch.export를 호출하여 PyTorch 모델로부터ExportedProgram(ATen Dialect)을 생성한다.- ExecuTorch의
to_edge()API를 사용하여 ATen Dialect 그래프를 엣지 환경에 더 적합한 Edge Dialect로 낮춘다. - 이후
to_executorch()API를 호출하여 최종적으로 엣지 디바이스에서 직접 실행 가능한.pte바이너리 파일을 생성한다.
7.2 ONNX: 상호운용성을 위한 새로운 표준
ONNX(Open Neural Network Exchange)는 프레임워크 간 모델 상호운용성을 위한 표준 형식이다. PyTorch 2.x에서는 torch.onnx.export(..., dynamo=True)를 사용하는 것이 새로운 표준 워크플로우로 권장된다.5 이 함수는 내부적으로 torch.export를 호출하여 건전한 FX 그래프를 캡처한 뒤, 이를 ONNX 형식으로 변환한다. 이는 TorchScript 기반의 레거시 ONNX 익스포터보다 훨씬 더 많은 PyTorch 모델을 안정적으로 변환할 수 있는 능력을 제공한다. 전체 워크플로우는 torch.export로 ExportedProgram을 얻고, 이를 ONNX로 변환한 후 ONNX Runtime과 같은 추론 엔진에서 실행하여 결과를 검증하는 과정을 포함한다.
7.3 TensorRT: 고성능 NVIDIA GPU 추론
NVIDIA의 고성능 추론 엔진인 TensorRT와 통합하기 위해, Torch-TensorRT 라이브러리는 ExportedProgram을 직접 입력으로 받아 TensorRT 엔진으로 컴파일하는 간소화된 경로를 제공한다.6
torch_tensorrt.dynamo.compile(exp_program, inputs)함수가 이 역할을 수행한다.- 이 함수는 내부적으로 그래프를 TensorRT가 지원하는 부분과 지원하지 않는 부분으로 분할하고, 지원되는 부분을 TensorRT 엔진으로 변환한 뒤, 이를 원래의 PyTorch 그래프 내에 임베드한 새로운
torch.fx.GraphModule을 반환한다. - 이 방식은 ONNX를 중간 형식으로 거치지 않고 PyTorch에서 TensorRT로 직접 변환하는 경로를 제공하여 워크플로우를 단순화한다.
ExportedProgram은 PyTorch 배포 생태계 내에서 “중심 허브(Central Hub)” 또는 “공용어(Lingua Franca)” 역할을 수행하도록 설계되었다. 과거에는 각 배포 타겟(ONNX, TensorRT, CoreML 등)마다 별도의, 종종 호환되지 않는 익스포트 경로가 존재하여 생태계가 파편화되어 있었다. 이제 모든 배포 워크플로우가 torch.export를 통해 생성된 표준화된 ExportedProgram에서 시작함으로써, 생태계 전체의 일관성과 통합성이 크게 향상되었다. 이는 새로운 배포 백엔드를 추가하는 작업을 극적으로 단순화시킨다. 새로운 백엔드 개발자는 PyTorch Eager Mode의 모든 복잡성을 이해할 필요 없이, 잘 정의된 ExportedProgram을 입력으로 받아 자신의 하드웨어에 맞게 낮추는(lowering) 로직만 구현하면 되기 때문이다. 결과적으로 torch.export는 PyTorch의 혁신 속도와 배포 생태계의 확장 속도를 동기화시키는 핵심적인 기술적 자산이 된다.
8. 결론: PyTorch 모델 배포의 미래
8.1 torch.export의 핵심 가치 요약
torch.export는 PyTorch 2.x의 강력한 컴파일러 기술을 활용하여, 연구 단계의 유연한 Python 모델과 프로덕션 환경의 고성능 정적 모델 사이의 깊은 간극을 메우는 가장 진보된 솔루션이다. Soundness, Normalization, Functionality라는 3대 원칙을 통해 생성된 ExportedProgram은 예측 가능하고, 이식 가능하며, 최적화에 용이한 표준 모델 표현을 제공함으로써, PyTorch 모델 배포의 새로운 시대를 열고 있다.
8.2 PyTorch 2.x 생태계에서의 전략적 위치와 미래 전망
torch.export는 torch.compile과 상호 보완적인 관계에 있다. torch.compile이 Python 환경 내에서의 학습 및 추론 성능 극대화를 추구한다면, torch.export는 Python을 벗어난 다양한 이기종 환경으로의 이식성을 책임진다. 이 두 축은 PyTorch 2.x가 지향하는 ’연구에서 프로덕션까지’의 엔드-투-엔드 경험을 완성하는 핵심 요소다.
앞으로 더 많은 하드웨어 벤더와 추론 엔진이 ExportedProgram을 직접 소비하는 방향으로 발전할 것이며, 이는 PyTorch 생태계의 파편화를 줄이고 통합을 가속화할 것이다. 동적 형상과 제어 흐름에 대한 지원이 계속 강화되면서, torch.export는 기존에 정적 그래프로 표현하기 어려웠던 Transformer, Graph Neural Network(GNN)와 같은 복잡한 모델의 배포를 더욱 용이하게 만들 것이다. 궁극적으로 torch.export는 PyTorch가 단순한 연구 프레임워크를 넘어, 진정한 엔드-투-엔드(end-to-end) AI 플랫폼으로 자리매김하는 데 있어 핵심적인 역할을 수행할 것이다.
9. 참고 자료
- torch.export — PyTorch 2.8 documentation, https://docs.pytorch.org/docs/stable/export.html
- Export IR Specification — ExecuTorch 0.7 documentation, https://docs.pytorch.org/executorch/stable/ir-exir.html
- ExportDB — PyTorch 2.8 documentation, https://docs.pytorch.org/docs/stable/generated/exportdb/index.html
- PyTorch Custom Operators — PyTorch Tutorials 2.8.0+cu128 …, https://docs.pytorch.org/tutorials/advanced/custom_ops_landing_page.html
- Export a PyTorch model to ONNX — PyTorch Tutorials 2.8.0+cu128 …, https://docs.pytorch.org/tutorials/beginner/onnx/export_simple_model_to_onnx_tutorial.html
- Compiling Exported Programs with Torch-TensorRT — Torch …, https://docs.pytorch.org/TensorRT/dynamo/dynamo_export.html